import numpy as np
from transforms3d.euler import mat2euler, euler2quat, quat2mat


def mat2quat(R):
    euler = mat2euler(R)
    q = euler2quat(euler[0], euler[1], euler[2])
    return q


def interp(T0, T1, alpha):  # alpha * T0 + (1 - alpha) * T1
    q0 = mat2quat(T0[:3, :3])
    t0 = T0[:3, 3]
    q1 = mat2quat(T1[:3, :3])
    t1 = T1[:3, 3]

    # interpolation
    if np.dot(q0, -q1) > np.dot(q0, q1):
        q1 = -q1
    q = alpha * q0 + (1 - alpha) * q1
    q /= np.linalg.norm(q, ord=2)
    t = alpha * t0 + (1 - alpha) * t1

    # return
    T = np.eye(4)
    T[:3, :3] = quat2mat(q)
    T[:3, 3] = t
    return T


def denoise_pose(rigid_pose, rigid_timestamps):

    assert len(rigid_pose) > 0
    
    valid_poses = [rigid_pose[0]]
    valid_ts = [rigid_timestamps[0]]

    last_valid = rigid_pose[0]  # 要求第0帧的pose必须是准的!!!
    last_valid_ts = rigid_timestamps[0]
    for i in range(1, len(rigid_pose)):
        if np.max(np.abs(rigid_pose[i] - rigid_pose[i-1])) < 1e-4:
            continue
        delta_deg = np.arccos(min(((np.trace(last_valid[:3, :3].T @ rigid_pose[i][:3, :3]) + 1) / 2) + 1e-6, 1.0)) / np.pi * 180
        av = delta_deg / ((rigid_timestamps[i] - last_valid_ts) * 1e-9)  # 角速率 (deg/s)
        v = np.linalg.norm(rigid_pose[i][:3, 3] - last_valid[:3, 3], ord=2) / ((rigid_timestamps[i] - last_valid_ts) * 1e-9)  # 线速率 (m/s)
        if (av > 180) or (v > 3):  # 速率过快则丢掉
            continue
        last_valid = rigid_pose[i]
        last_valid_ts = rigid_timestamps[i]
        valid_poses.append(rigid_pose[i])
        valid_ts.append(rigid_timestamps[i])
        
    
    print(len(rigid_pose), len(valid_poses))

    return valid_poses, valid_ts
